%% 无人机模型训练与评估主程序
clear; clc; close all force;

%% 0. 系统诊断检查
disp('=== 系统初始化检查 ===');
% 检查图形系统
try
    testFig = figure('Visible','off');
    close(testFig);
    disp('图形系统: 正常');
catch
    warning('图形系统: 异常');
end

% 检查Python连接
try
    py.importlib.import_module('sys');
    disp('Python连接: 正常');
catch
    error('Python连接失败，请检查Python环境');
end

%% 1. 路径配置
pythonModulePath = 'D:\Pycharmcode\test2\test2';
dataFolder = 'D:\Pycharmcode\test2\data';
modelSavePath = 'D:\Pycharmcode\test2\matlabtrain\drone_model_trained.pth';

% 检查数据路径
assert(isfolder(dataFolder), '数据路径不存在: %s', dataFolder);
disp('数据路径验证: 通过');

%% 2. 系统初始化
% 设置Python环境
pe = pyenv;
if pe.Status == "NotLoaded"
    pyenv('Version', 'D:\Anaconda3\envs\neural_fly\python.exe');
end

% 添加Python路径
if count(py.sys.path, pythonModulePath) == 0
    insert(py.sys.path, int32(0), pythonModulePath);
end

% 验证模块导入
try
    py.importlib.import_module('drone_model');
    disp('Python模块导入成功');
    
    % 获取Python配置
    py_config = py.drone_model.Config();
    disp('Python配置参数:');
    disp(py_config);
catch e
    error('Python模块导入失败: %s', e.message);
end

% 创建模型保存目录
[modelSaveDir, ~, ~] = fileparts(modelSavePath);
if ~isfolder(modelSaveDir)
    mkdir(modelSaveDir);
end

%% 3. 准备数据
try
    [trainData, testData] = prepareData(dataFolder);
    
    fprintf('\n=== 数据加载结果 ===\n');
    fprintf('训练样本数: %d\n', numel(trainData));
    fprintf('测试样本数: %d\n', numel(testData));
    
    % 显示数据统计信息
    fprintf('\n训练数据分布:\n');
    conditions = cellfun(@(x) x.condition, trainData, 'UniformOutput', false);
    tabulate(conditions);
    
    % 检查数据维度
    fprintf('\n首个训练样本维度检查:\n');
    sample = trainData{1};
    disp(struct(...
        'acc_size', size(sample.acc), ...
        'vel_size', size(sample.vel), ...
        'gyro_size', size(sample.gyro), ...
        'pos_size', size(sample.pos), ...
        'fa_size', size(sample.fa), ...
        'features_size', size(sample.features)));
    
catch ME
    error('数据加载失败: %s', ME.message);
end

%% 4. 创建并训练模型
try
    % 创建模型实例
    model = DroneModel(py_config);
    fprintf('\n模型特征配置: %s\n', strjoin(model.feature_names, ', '));
    disp('模型初始化成功');
    
    % 训练模型
    fprintf('\n=== 开始训练模型 ===\n');
    tic;
    model.train(trainData);
    trainingTime = toc;
    fprintf('训练完成，耗时: %.2f秒\n', trainingTime);
    
    % 保存模型
    model.saveModel(modelSavePath);
    fprintf('模型已保存至: %s\n', modelSavePath);
    
catch ME
    error('模型训练失败: %s\n%s', ME.message, getReport(ME, 'extended'));
end

%% 5. 测试预测
try
    disp('开始测试预测...');
    
    % 使用第一个测试样本
    testSample = testData{1};
    
    % 准备测试输入 (使用完整features)
    testInput = testSample.features(1:5,:); % 取前5个时间步
    disp(['输入数据维度: ', mat2str(size(testInput))]);
    
    % 准备适配数据 (同样使用features)
    adaptSample = trainData{1};
    adaptX = adaptSample.features(1:32,:); % 取32个样本
    adaptY = adaptSample.fa(1:32,:);
    
    % 执行预测
    tic;
    result = model.predict(testInput, {adaptX, adaptY});
    predictionTime = toc;
    
    %% 输出结果
    disp('===== 预测结果 =====');
    disp('预测气动力(N):');
    disp(array2table(result.forces, 'VariableNames', {'X','Y','Z'}));
    
    disp('预测风况:');
    if iscell(result.wind_conditions)
        disp(strjoin(result.wind_conditions, ', '));
    else
        disp(result.wind_conditions);
    end
    
    fprintf('预测耗时: %.4f秒\n', predictionTime);
    
catch ME
    warning(ME.identifier,'预测错误: %s', ME.message);
    disp('当前变量维度:');
    whos testInput adaptX adaptY;
end

%% 6. 完整模型评估
try
    fprintf('\n=== 开始模型评估 ===\n');
    
    % 评估模型性能
    [model_perf, forces_pred, forces_true] = evaluate_model(model, testData);
    
    % 风况独立性分析
    [wind_indep, feat_analysis] = evaluate_wind_independence(model, testData);
    
    % 适配能力评估
    adapt_curve = evaluate_adaptation(model, trainData, testData);
    
    % 可视化结果
    visualize_performance(model_perf, wind_indep, feat_analysis, adapt_curve,...
        forces_pred, forces_true, testData);
    
catch ME
    warning(ME.identifier,'评估过程中出错: %s', ME.message);
end

%% 7. 辅助函数定义
function [model_perf, forces_pred, forces_true] = evaluate_model(model, testData)
    % 初始化
    samples_per_case = size(testData{1}.fa, 1);
    total_samples = numel(testData) * samples_per_case;
    forces_pred = zeros(total_samples, 3);
    forces_true = zeros(total_samples, 3);
    
    % 批量预测（使用完整features）
    for i = 1:numel(testData)
        idx = (i-1)*samples_per_case+1 : i*samples_per_case;
        pred = model.predict(testData{i}.features);
        forces_pred(idx,:) = pred.forces;
        forces_true(idx,:) = testData{i}.fa;
    end
    
    % 计算误差指标
    abs_error = mean(abs(forces_pred - forces_true));
    rel_error = abs_error ./ mean(abs(forces_true));
    angle_error = acosd(dot(forces_pred, forces_true, 2))./...
        (vecnorm(forces_pred,2,2).*vecnorm(forces_true,2,2));
    
    model_perf = struct(...
        'MAE', abs_error,...
        'MRE', rel_error,...
        'RMSE', sqrt(mean((forces_pred - forces_true).^2)),...
        'MaxError', max(abs(forces_pred - forces_true)),...
        'Correlation', diag(corr(forces_pred, forces_true))',...
        'AngleError', mean(angle_error),...
        'StdRatio', std(forces_pred - forces_true) ./ std(forces_true));
    
    fprintf('评估结果:\n平均绝对误差: %.4f N\n角度误差: %.1f°\n', ...
        mean(model_perf.MAE), model_perf.AngleError);
end

function [wind_indep, feat_analysis] = evaluate_wind_independence(model, testData)
    % 特征提取
    fprintf('正在提取特征...\n');
    features = cellfun(@(x) model.py_model.get_phi_features(...
        py.numpy.array(x.features)), testData, 'UniformOutput', false);
    features = cat(1, features{:});
    
    conditions = cellfun(@(x) x.condition, testData, 'UniformOutput', false);
    
    % PCA分析
    fprintf('正在进行PCA分析...\n');
    [coeff, score] = pca(double(features));
    pca_var = cumsum(var(score))/sum(var(score));
    
    % 方差分析
    [~,~,stats] = anova1(score(:,1), conditions, 'off');
    
    wind_indep = struct(...
        'ANOVA_p', stats.p,...
        'PCA_variance', pca_var(1:3),...
        'Cluster_silhouette', mean(silhouette(score(:,1:3), grp2idx(conditions))));
    
    feat_analysis = struct(...
        'PCA_coeff', coeff(:,1:3),...
        'PCA_scores', score(:,1:3),...
        'Conditions', {conditions});
end

function adapt_curve = evaluate_adaptation(model, trainData, testData)
    adapt_steps = min(10, numel(trainData));
    errors = zeros(adapt_steps,1);
    test_sample = testData{1};
    
    fprintf('评估适配能力(%d步)...\n', adapt_steps);
    
    for k = 1:adapt_steps
        % 渐进式适配
        adapt_X = trainData{k}.features(1:32,:);
        adapt_Y = trainData{k}.fa(1:32,:);
        model.py_model.adapt(py.numpy.array(adapt_X), py.numpy.array(adapt_Y));
        
        % 测试当前适配效果
        pred = model.predict(test_sample.features);
        errors(k) = mean(vecnorm(pred.forces - test_sample.fa, 2, 2));
    end
    
    adapt_curve = struct(...
        'samples', 1:adapt_steps,...
        'errors', errors,...
        'final_error', errors(end));
    
    fprintf('最终适配误差: %.4f N\n', errors(end));
end

function visualize_performance(model_perf, wind_indep, feat_analysis, adapt_curve,...
    forces_pred, forces_true, testData)
    
    % 创建图形窗口
    fig = figure('Name','无人机模型性能评估', ...
        'NumberTitle','off', ...
        'Position',[100 100 1400 900], ...
        'Color','w');
    
    try
        %% 1. 误差指标对比
        subplot(2,3,1);
        barh([model_perf.MAE; model_perf.RMSE]');
        set(gca,'YTickLabel',{'X轴','Y轴','Z轴'});
        title('力预测误差对比 (N)');
        legend({'MAE','RMSE'},'Location','best');
        grid on;
        
        %% 2. 风况独立性分析
        subplot(2,3,2);
        boxplot(feat_analysis.PCA_scores(:,1), feat_analysis.Conditions);
        xtickangle(45);
        title(sprintf('风况独立性 (p=%.2e)', wind_indep.ANOVA_p));
        ylabel('第一主成分值');
        grid on;
        
        %% 3. 动态适应曲线
        subplot(2,3,3);
        plot(adapt_curve.samples, adapt_curve.errors, '-o', 'LineWidth',1.5);
        xlabel('适配样本数');
        ylabel('误差(N)');
        title(sprintf('动态适应 (最终误差: %.3f N)', adapt_curve.final_error));
        grid on;
        
        %% 4. 三维误差空间
        subplot(2,3,4);
        scatter3(forces_true(:,1), forces_true(:,2), forces_true(:,3), ...
            50, 'b', 'filled', 'MarkerEdgeColor','k');
        hold on;
        scatter3(forces_pred(:,1), forces_pred(:,2), forces_pred(:,3), ...
            30, 'r', 'filled');
        legend('真实值','预测值', 'Location','best');
        xlabel('X轴 (N)'); ylabel('Y轴 (N)'); zlabel('Z轴 (N)');
        title('三维力空间分布');
        grid on; axis equal;
        
        %% 5. 误差角度分布
        subplot(2,3,5);
        polarhistogram(acosd(dot(forces_pred, forces_true, 2)./...
            (vecnorm(forces_pred,2,2).*vecnorm(forces_true,2,2))), 20, ...
            'FaceColor','m', 'EdgeColor','k');
        title(sprintf('力向量角度误差 (平均: %.1f°)', model_perf.AngleError));
        
        %% 6. 性能指标表
        subplot(2,3,6);
        axis off;
        metrics = {
            '平均绝对误差 (N)', sprintf('%.3f ± %.3f', mean(model_perf.MAE), std(model_perf.MAE));
            '最大误差 (N)', sprintf('%.3f', max(model_perf.MaxError));
            '角度误差 (°)', sprintf('%.1f ± %.1f', model_perf.AngleError, std(acosd(dot(forces_pred, forces_true, 2))./...
            (vecnorm(forces_pred,2,2).*vecnorm(forces_true,2,2))));
            '风况p值', sprintf('%.2e', wind_indep.ANOVA_p);
            'PCA解释率', sprintf('%.1f%%', wind_indep.PCA_variance(3)*100);
            '适配后误差 (N)', sprintf('%.3f', adapt_curve.final_error)};
        
        uitable(fig, 'Data', metrics, ...
            'Position', [700 100 600 200], ...
            'ColumnName', {'指标', '值'}, ...
            'RowName', [], ...
            'FontSize', 10);
        
        drawnow;
    catch ME
        warning(ME.identifier,'可视化错误: %s', ME.message);
        if isvalid(fig)
            close(fig);
        end
    end
end

%% 8. 脚本完成
disp('===== 脚本执行完成 =====');
fprintf('完成时间: %s\n', string(datetime("now")));
